import numpy as np
import matplotlib as mpl
# 使用sklearn随机生成数据
# make_classification生成三元分类模型数据
from sklearn.datasets import make_classification
import matplotlib.pyplot as plt

"""
极限学习机的一个示例, 可以正确运行
https://blog.csdn.net/weixin_41534404/article/details/104819523?ops_request_misc=%257B%2522request%255Fid%2
522%253A%2522160214554519724836717618%2522%252C%2522scm%2522%253A%252220140713.130102334..%2522%257D&reques
t_id=160214554519724836717618&biz_id=0&utm_medium=distribute.pc_search_result.no
ne-task-blog-2~all~first_rank_v2~rank_v28_p-2-104819523.pc_first_rank_v2_rank_v28_p&utm_term=ELM+pytorch
"""

class ELM():
    # 输入数据集X、标签Y、神经元个数m、控制参数L
    def __init__(self, X, Y, m, L):
        self.X = X
        self.Y = Y
        self.m, self.L = m, L
        self.TRAIN_beta()

    def sigmoid(self, x):
        return 1.0 / (1 + np.exp(-x))

    # 训练函数，随机w,b 计算H、beta
    def TRAIN_beta(self):
        n, d = self.X.shape
        self.w = np.random.rand(d, self.m)
        self.b = np.random.rand(1, self.m)
        H = self.sigmoid(np.dot(self.X, self.w) + self.b)
        self.beta = np.dot(np.linalg.inv(np.identity(self.m) / self.L + np.dot(H.T, H)),
                           np.dot(H.T, self.Y))  # 加入正则化且 n >> m
        # self.beta = np.dot(np.linalg.pinv(H), self.Y) # 不加入正则化
        print('TRAIN FINISH beta_shape ', self.beta.shape)

    # 测试函数，计算方法
    def TEST(self, x):
        H = self.sigmoid(np.dot(x, self.w) + self.b)  # 使用测试集计算H，其中w、b是之前随机得到的
        result = np.dot(H, self.beta)
        return result



# 关键参数有n_samples（生成样本数）， n_features（样本特征数）， n_redundant（冗余特征数）和n_classes（输出的类别数）
# X1为样本特征，Y1为样本类别输出， 共400个样本，每个样本2个特征，输出有2个类别，没有冗余特征，每个类别一个簇
X,Y = make_classification(n_samples=10000, n_features=2, n_redundant=0, n_clusters_per_class=1, n_classes=3)
plt.scatter(X[:,0], X[:,1], c=Y, s=3, marker='o')
plt.show()


# 训练时将标签改为OneHot编码效果比较好
Y_onehot = np.eye(3)[Y]

# 用前8000个数据训练，5000个神经元，控制参数取0.1
elm = ELM(X[:8000], Y_onehot[:8000], 5000, 0.1)
# 用后2000个数据进行测试
predict = elm.TEST(X[8000:])
predict = np.argmax(predict, axis=1) # OneHot编码形式 取每行最大值的索引即类别
acc = np.sum(predict == Y[8000:])/2000
print('acc :', acc)

# 绘制分类彩图

x_np = np.asarray(X)
y_np = np.asarray(Y)
N, M = 50, 50 # 横纵各采样多少个值

x1_min, x2_min = x_np.min(axis=0) # 列最小值
x1_max, x2_max = x_np.max(axis=0) # 列最大值

t1 = np.linspace(x1_min, x1_max, N) # 在min-max间线性采样
t2 = np.linspace(x2_min, x2_max, M)

x1, x2 = np.meshgrid(t1, t2) # 生成网格采样点
x_show = np.stack((x1.flat, x2.flat), axis=1) # 测试点
y_predict = np.argmax(elm.TEST(x_show), axis=1)

# 绘制分类图
cm_light = mpl.colors.ListedColormap(['#A0FFA0', '#FFA0A0', '#A0A0FF'])
cm_dark = mpl.colors.ListedColormap(['g', 'r', 'b'])
plt.pcolormesh(x1, x2, np.array(y_predict).reshape(x1.shape), cmap=cm_light)
plt.scatter(X[:, 0], X[:, 1], c=Y, s=3, marker='o')
plt.show()










